import json
import random
import websocket
import threading
import time
import sys
from colorama import init, AnsiToWin32
from blessed import Terminal


class WebSSHClient:
    def __init__(self, url, auth):
        init()
        self.ws_url = url
        self.auth = auth
        self.ws = None
        self.screen = ''
        self.exit_status = -1
        self.running = False
        self.script_output = False
        self._zmodem_buffer = bytearray()
        self.ansi_converter = AnsiToWin32(sys.stdout)
        self.term = Terminal()
        self._connection_complete = threading.Event()  # 连接完成标志

    def _connect_thread(self):
        """在独立线程中执行连接"""
        try:
            self.ws = websocket.WebSocketApp(
                self.ws_url,
                on_open=self.on_open,
                on_message=self.on_message,
                on_error=self.on_error,
                on_close=self.on_close,
            )
            self.ws.run_forever()
        except Exception as e:
            print(f"Connection failed: {e}")
            self.running = False

    def _send_normal_key(self, key):
        """发送普通按键"""
        command = {
            "operate": "COMMAND",
            "device": "HOST",
            "connect": "SSH",
            "command": str(key)
        }
        self.ws.send(json.dumps(command))

    def connect(self):
        """非阻塞启动连接"""
        self.running = True
        conn_thread = threading.Thread(target=self._connect_thread, daemon=True)
        conn_thread.start()

        # 等待连接完成或失败
        while self.running:
            if self._connection_complete.wait(timeout=0.1):
                break
            if not conn_thread.is_alive():
                break

    def on_message(self, ws, message):
        pass

    def on_error(self, ws, error):
        print(f"\nWebSocket error: {error}")
        self.running = False

    def on_close(self, ws, close_status_code, close_msg):
        self.running = False

    def on_open(self, ws):
        """连接成功回调"""
        self._connection_complete.set()
        ws.send(json.dumps(self.auth))

    def run_script(self, content, interpreter):
        _wait = 0
        while _wait < 60:
            if not self.screen.strip().endswith('~]$'):
                time.sleep(0.5)
                _wait += 1
                continue

            random_number = random.randint(10, 100)
            filename = f'ops_{random_number}.py'
            self.filename = filename
            new_content = f"""
cat > /tmp/{filename} << 'EOF'
{content}
print()
EOF
"""
            self._send_normal_key(new_content)
            self._send_normal_key(f'{interpreter} /tmp/{filename} # ---start--- \n')
            self._send_normal_key(f'echo $? # ---end--- \n')
            self._send_normal_key(f'rm -f /tmp/{filename}  \n')
            break


def run_script_yunqing(content: str):
    ws_url = "ws://172.17.42.211:9090/wb/socket/gy/ws"
    auth = {
        "certId": 12328,
        "port": 22,
        "host": "10.126.8.46",
        "hostId": "121wi7POOy",
        "encryptCredential": "T9PhHVPWRGhdNfYZEmN/qwdwupmL5JH0fhlOL+odTlw=",
        "operate": "CONNECT",
        "device": "HOST",
        "connect": "SSH",
        "col": 213,
        "row": 21
    }
    interpreter = 'cd /data/aops/aops &&  sudo /usr/local/python39/bin/python3'

    client = WebSSHClient(ws_url, auth)

    # 创建一个队列用于存储收到的消息
    message_queue = []

    def on_message(ws, message):
        client.screen += message
        if '---start---' in message:
            client.script_output = True
            return
        elif '---end---' in message:
            print('end: ', message)
            client.script_output = False
            exit_status = message.strip()[-1]
            if exit_status.isdigit():
                client.exit_status = exit_status
            else:
                client.exit_status = 0
            return

        if client.script_output:
            message_queue.append(message)

    client.on_message = on_message

    client.connect()
    client.run_script(content, interpreter)

    while True:
        if '---end---' in client.screen:
            break
        # 如果有新消息，就yield它们
        while message_queue:
            yield message_queue.pop(0)
        time.sleep(0.1)

    # yield剩余的消息
    while message_queue:
        yield message_queue.pop(0)
    yield {
        'end': True,
        'exit_status': client.exit_status,
    }

    client.ws.close()
